
import org.tensorflow.Graph;
import org.tensorflow.Session;
import org.tensorflow.Tensor;

import java.io.BufferedReader;
import java.io.FileInputStream;
import java.io.IOException;
import java.io.InputStreamReader;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.*;


public class BertNerPredict {
    private static String vocabPath = "src/main/resources/vocab.txt";
    private static Map<String, Integer> word2id = new HashMap<String, Integer>();
    static {
        try {
            BufferedReader buffer = null;
            buffer = new BufferedReader(new InputStreamReader(new FileInputStream(vocabPath)));
            int i = 0;
            String line = buffer.readLine().trim();
            while (line!=null){
                word2id.put(line, i++);
                line = buffer.readLine().trim();
            }
            buffer.close();
        }catch (Exception e){
        }
//        System.out.println("word2id size is:"+word2id.size());

    }

    private static byte[] readAllByteOrExit(Path path){
        try{
            return Files.readAllBytes(path);
        }catch (IOException e){
            System.out.println("Failed to read[" + path + "]:" + e.getMessage());
            System.exit(1);
        }
        return null;
    }

    public static void getTextToId(int[][] inputIds, int[][] inputMask, String[] text){
        for(int i=0; i<text.length; i++){
            char[] chs = text[i].trim().toLowerCase().toCharArray();

            List<Integer> list = new ArrayList<>();
            List<Integer> mask = new ArrayList<>();
            list.add(word2id.get("[CLS]"));
            mask.add(1);
            for(int j=0; j<chs.length; j++){
                String element = Character.toString(chs[j]);
                if(word2id.containsKey(element)){
                    list.add(word2id.get(element));
                    mask.add(1);
                }
            }
            list.add(word2id.get("[SEP]"));
            mask.add(1);

            int size = list.size();
            Integer[] targetInter = list.toArray(new Integer[size]);
            int[] idResult = Arrays.stream(targetInter).mapToInt(Integer::valueOf).toArray();
            Integer[] maskInter = mask.toArray(new Integer[size]);
            int[] maskResult = Arrays.stream(maskInter).mapToInt(Integer::valueOf).toArray();
            System.arraycopy(idResult,0,inputIds[i], 0, idResult.length);
            System.arraycopy(maskResult,0,inputMask[i], 0, maskResult.length);
        }
    }

    public static void main(String[] args) {
        String[] query = new String[]{"中华人民共和国", "新疆大学"};




        String resourceDir = "src/main/resources";
        String modelName = "model.pb";

        int batchSize = query.length;
        int padLength = 25;
        int[][] indexes = new int[batchSize][padLength];
        int[][] mask = new int[batchSize][padLength];

        byte[] graphDef = readAllByteOrExit(Paths.get(resourceDir, modelName));
        Graph g = new Graph();
        g.importGraphDef(graphDef);
        Session sess = new Session(g);

        if (query.length>0){
            System.out.println("Ok! Start predicting...\n");
        }else {
            System.exit(0);
        }

        getTextToId(indexes, mask, query);
        Tensor tensorInputIds = Tensor.create(indexes);
        Tensor tensorMask = Tensor.create(mask);
        Tensor tensorSeqResult = sess.runner().feed("input_ids:0", tensorInputIds).
                feed("input_mask:0", tensorMask).fetch("viterbi/ReverseSequence_1:0").run().get(0);
        Tensor tensorScoreResult = sess.runner().feed("input_ids:0", tensorInputIds).
                feed("input_mask:0", tensorMask).fetch("viterbi/Max:0").run().get(0);
        int[][] sequenceId = (int[][]) tensorSeqResult.copyTo(new int[batchSize][padLength]);
        float[] sequenceScore = (float[]) tensorScoreResult.copyTo(new float[batchSize]);
        for(int i=0; i<sequenceId.length; i++){
            System.out.println("query: "+query[i]);
            System.out.println("sequence result: "+ Arrays.toString(sequenceId[i]));
            System.out.println("sequence score: "+ sequenceScore[i]+'\n');
        }
        tensorScoreResult.close();
        tensorSeqResult.close();
        tensorMask.close();
        tensorInputIds.close();
    }
}